#include <iostream>
#include <string>
#include <vector>
#include <ctime>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include "task.hpp"

using namespace std;

enum
{
    UsageError = 1,
    ArgError,
    PipeError
};

class Channel
{
public:
    Channel(int wfd, pid_t sub_id, const string &name)
        : _wfd(wfd), _sub_process_pid(sub_id), _name(name)
    {
    }

    void PrintDebug()
    {
        cout << "wfd: " << _wfd << " pid: " << _sub_process_pid << " name: " << _name << endl;
    }

    void Close() { close(_wfd); }
    int wfd() const { return _wfd; }
    pid_t pid() const { return _sub_process_pid; }
    string name() const { return _name; }

    ~Channel()
    {
    }

private:
    int _wfd;
    pid_t _sub_process_pid;
    string _name;
};

class ProcessPool
{
public:
    ProcessPool(int sub_process_num) : _sub_process_num(sub_process_num)
    {
    }

    int CreateProcess(work_t work) // 回调函数
    {
        vector<int> fds;
        for (int number = 0; number < _sub_process_num; number++)
        {
            int pipefd[2] = {0};
            int n = pipe(pipefd);
            if (n == -1)
            {
                perror("pipe error");
                return PipeError;
            }

            pid_t id = fork();
            if (id == 0)
            {
                if (!fds.empty())
                {
                    cout << "close w fd:";
                    for (auto fd : fds)
                    {
                        close(fd);
                        cout << fd << " ";
                    }
                    cout << endl;
                }

                // child -> r
                close(pipefd[1]);

                // 执行任务
                dup2(pipefd[0], STDIN_FILENO);
                work(pipefd[0]);
                
                //cout<<"child exit, pid:"<<getpid()<<endl;
                exit(0);
            }
            string name = "channel-" + to_string(number);

            // father
            close(pipefd[0]);
            channels.push_back(Channel(pipefd[1], id, name));

            // 保存父进程的wfd
            fds.push_back(pipefd[1]);
        }
        return 0;
    }

    int NextChannel()
    {
        static int index = 0;
        int cpy = index;
        index++;
        index %= channels.size();
        return cpy;
    }

    void SendTaskCode(int channel_index, uint32_t code)
    {
        cout << "send code: " << code << " to " << channels[channel_index].name()
             << ", sub process id: " << channels[channel_index].pid() << endl;
        write(channels[channel_index].wfd(), &code, sizeof(code));
    }

    void KillAll()
    {
        for (auto &channel : channels)
        {
            channel.Close();
            pid_t pid = channel.pid();

            pid_t rid = waitpid(pid, NULL, 0);
            if (rid == pid)
            {
                cout << "wait sub processs " << pid << " success..." << endl;
            }
            cout << channel.name() << " close done, sub process quit now : " << channel.pid() << endl;
        }
    }

    void Debug()
    {
        for (auto &channel : channels)
        {
            channel.PrintDebug();
        }
    }

    ~ProcessPool()
    {
    }

private:
    int _sub_process_num;
    vector<Channel> channels;
};

void CtrlProcessPool(ProcessPool *processpool_ptr, int cnt)
{
    while (cnt)
    {
        // 选择一个进程和通道
        int channel = processpool_ptr->NextChannel();

        // 选择一个任务
        uint32_t code = NextTaskCode();

        // 发送任务
        processpool_ptr->SendTaskCode(channel, code);

        sleep(1);
        cnt--;
    }
}

int main(int argc, char *argv[])
{
    if (argc != 2)
    {
        cout << "Usage: " << argv[0] << " subprocess-num" << endl;
        return UsageError;
    }
    int sub_process_num = stoi(argv[1]);
    if (sub_process_num <= 0)
        return ArgError;

    srand((uint64_t)time(NULL));

    // 创建通信信道和子进程
    ProcessPool *processpool_ptr = new ProcessPool(sub_process_num);
    processpool_ptr->CreateProcess(worker);
    processpool_ptr->Debug();

    // 控制子进程
    CtrlProcessPool(processpool_ptr, 10);

    std::cout << "task run done" << std::endl;

    // 回收子进程：退出并等待
    processpool_ptr->KillAll();

    delete processpool_ptr;
    return 0;
}
